#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
知识搜索功能修复验证测试

本测试文件严格遵循以下测试规范：
1. 直接调用项目源码暴露的接口进行验证
2. 测试知识上传、向量化、搜索的完整流程
3. 验证向量恢复功能是否正常工作
4. 检查数据库连接、查询逻辑、返回结果处理
5. 发现问题时直接修正源码缺陷
"""

import sys
import os
import asyncio
import logging
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))
sys.path.insert(0, str(project_root / "src"))

# 设置环境变量
os.environ['PYTHONPATH'] = str(project_root / "src")

# 导入项目核心模块
from src.core.knowledge_processor import KnowledgeProcessor
from src.database.manager import DatabaseManager
from src.config import load_config

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class 知识搜索功能修复验证测试:
    """知识搜索功能修复验证测试类"""
    
    def __init__(self):
        self.config = None
        self.db_manager = None
        self.knowledge_processor = None
        self.test_results = []
        
    async def setup(self):
        """测试环境初始化"""
        try:
            # 加载配置
            self.config = load_config()
            logger.info("配置加载成功")
            
            # 初始化数据库管理器
            self.db_manager = DatabaseManager(self.config)
            await self.db_manager.initialize()
            logger.info("数据库管理器初始化成功")
            
            # 初始化知识处理器
            try:
                self.knowledge_processor = KnowledgeProcessor(self.db_manager, self.config)
                await self.knowledge_processor.initialize()
                logger.info("知识处理器初始化成功")
            except Exception as ke:
                logger.error(f"知识处理器初始化失败: {ke}")
                import traceback
                logger.error(f"详细错误信息: {traceback.format_exc()}")
                raise ke
            
            return True
            
        except Exception as e:
            logger.error(f"测试环境初始化失败: {e}")
            return False
    
    async def test_系统初始化(self):
        """测试系统初始化功能"""
        logger.info("开始测试系统初始化功能")
        
        try:
            # 验证配置对象
            assert self.config is not None, "配置对象未正确初始化"
            
            # 验证数据库管理器
            assert self.db_manager is not None, "数据库管理器未正确初始化"
            
            # 验证知识处理器
            assert self.knowledge_processor is not None, "知识处理器未正确初始化"
            
            # 检查关键组件是否正确初始化
            assert hasattr(self.knowledge_processor, 'vectorizer'), "向量化器未初始化"
            assert self.knowledge_processor.vectorizer is not None, "向量化器为None"
            
            # 检查向量恢复功能
            if (hasattr(self.knowledge_processor, 'vectorizer') and 
                self.knowledge_processor.vectorizer and 
                hasattr(self.knowledge_processor.vectorizer, 'vector_store')):
                vector_store = self.knowledge_processor.vectorizer.vector_store
                vector_count = len(vector_store.vectors) if hasattr(vector_store, 'vectors') else 0
                logger.info(f"当前向量数量: {vector_count}")
            else:
                logger.error("向量存储不可用")
                assert False, "向量存储未正确初始化"
            
            self.test_results.append(("系统初始化测试", "通过", "所有核心组件正确初始化"))
            logger.info("系统初始化测试通过")
            return True
            
        except Exception as e:
            self.test_results.append(("系统初始化测试", "失败", str(e)))
            logger.error(f"系统初始化测试失败: {e}")
            return False
    
    async def test_数据库连接状态(self):
        """测试数据库连接状态"""
        logger.info("开始测试数据库连接状态")
        
        try:
            # 测试ChromaDB连接
            chroma_client = self.db_manager.chromadb_client
            assert chroma_client is not None, "ChromaDB客户端未初始化"
            
            # 测试MongoDB连接
            mongo_client = self.db_manager.mongodb_client
            assert mongo_client is not None, "MongoDB客户端未初始化"
            
            # 测试Neo4j连接
            neo4j_client = self.db_manager.neo4j_client
            assert neo4j_client is not None, "Neo4j客户端未初始化"
            
            self.test_results.append(("数据库连接测试", "通过", "所有数据库连接正常"))
            logger.info("数据库连接测试通过")
            return True
            
        except Exception as e:
            self.test_results.append(("数据库连接测试", "失败", str(e)))
            logger.error(f"数据库连接测试失败: {e}")
            return False
    
    async def test_文档上传功能(self):
        """测试文档上传和向量化功能"""
        logger.info("开始测试文档上传功能")
        
        try:
            # 准备测试文档
            test_content = "这是一个测试文档，用于验证知识搜索功能修复。文档包含关于MCP-2系统的重要信息。"
            test_metadata = {
                "title": "测试文档",
                "source": "测试",
                "type": "验证文档"
            }
            
            # 调用知识上传接口
            result = await self.knowledge_processor.upload_document(
                content=test_content,
                title=test_metadata["title"],
                metadata=test_metadata
            )
            
            # 验证上传结果
            assert result is not None, "文档上传返回结果为None"
            assert "success" in result, "上传结果缺少success字段"
            assert result["success"] is True, f"文档上传失败: {result.get('error', '未知错误')}"
            
            # 验证文档ID
            assert "document_id" in result, "上传结果缺少document_id"
            document_id = result["document_id"]
            assert document_id is not None, "文档ID为None"
            
            self.test_results.append(("文档上传测试", "通过", f"文档ID: {document_id}"))
            logger.info(f"文档上传测试通过，文档ID: {document_id}")
            return document_id
            
        except Exception as e:
            self.test_results.append(("文档上传测试", "失败", str(e)))
            logger.error(f"文档上传测试失败: {e}")
            return None
    
    async def test_知识搜索功能(self, test_query="MCP-2系统"):
        """测试知识搜索功能"""
        logger.info(f"开始测试知识搜索功能，查询: {test_query}")
        
        try:
            # 调用知识搜索接口
            search_results = await self.knowledge_processor.search_knowledge(
                query=test_query,
                limit=5
            )
            
            # 验证搜索结果结构
            assert search_results is not None, "搜索结果为None"
            assert isinstance(search_results, list), f"搜索结果类型错误: {type(search_results)}"
            
            # 验证搜索结果内容
            if len(search_results) > 0:
                # 检查结果项结构
                for result in search_results:
                    assert isinstance(result, dict), "搜索结果项不是字典类型"
                    assert "content" in result, "搜索结果缺少content字段"
                    assert "score" in result, "搜索结果缺少score字段"
                    assert "metadata" in result, "搜索结果缺少metadata字段"
                
                self.test_results.append(("知识搜索测试", "通过", f"找到{len(search_results)}个结果"))
                logger.info(f"知识搜索测试通过，找到{len(search_results)}个结果")
                return search_results
            else:
                # 搜索结果为空，可能是向量恢复问题
                logger.warning("搜索结果为空，可能需要检查向量恢复功能")
                self.test_results.append(("知识搜索测试", "警告", "搜索结果为空"))
                return []
                
        except Exception as e:
            self.test_results.append(("知识搜索测试", "失败", str(e)))
            logger.error(f"知识搜索测试失败: {e}")
            return None
    
    async def test_向量恢复功能(self):
        """测试向量恢复功能"""
        logger.info("开始测试向量恢复功能")
        
        try:
            # 检查是否有_restore_vectors_from_chromadb方法
            assert hasattr(self.knowledge_processor, '_restore_vectors_from_chromadb'), \
                "知识处理器缺少_restore_vectors_from_chromadb方法"
            
            # 获取当前向量存储状态
            vector_store = self.knowledge_processor.vectorizer.vector_store
            initial_count = len(vector_store.vectors) if hasattr(vector_store, 'vectors') else 0
            
            logger.info(f"向量恢复前内存中向量数量: {initial_count}")
            
            # 调用向量恢复方法
            await self.knowledge_processor._restore_vectors_from_chromadb()
            
            # 检查恢复后的状态
            final_count = len(vector_store.vectors) if hasattr(vector_store, 'vectors') else 0
            logger.info(f"向量恢复后内存中向量数量: {final_count}")
            
            # 验证向量恢复效果
            if final_count >= initial_count:
                self.test_results.append(("向量恢复测试", "通过", f"恢复了{final_count - initial_count}个向量"))
                logger.info(f"向量恢复测试通过，恢复了{final_count - initial_count}个向量")
                return True
            else:
                self.test_results.append(("向量恢复测试", "警告", "向量数量未增加"))
                logger.warning("向量恢复测试警告：向量数量未增加")
                return False
                
        except Exception as e:
            self.test_results.append(("向量恢复测试", "失败", str(e)))
            logger.error(f"向量恢复测试失败: {e}")
            return False
    
    async def test_错误处理机制(self):
        """测试错误处理机制"""
        logger.info("开始测试错误处理机制")
        
        try:
            # 测试无效查询处理
            invalid_results = await self.knowledge_processor.search_knowledge(
                query="",  # 空查询
                limit=5
            )
            
            # 验证空查询的处理
            assert invalid_results is not None, "空查询应返回空列表而不是None"
            assert isinstance(invalid_results, list), "空查询结果应为列表类型"
            
            # 测试超大限制值处理
            large_limit_results = await self.knowledge_processor.search_knowledge(
                query="测试",
                limit=10000  # 超大限制值
            )
            
            # 验证超大限制值的处理
            assert large_limit_results is not None, "超大限制值查询应正常处理"
            assert isinstance(large_limit_results, list), "超大限制值查询结果应为列表类型"
            
            self.test_results.append(("错误处理测试", "通过", "异常情况处理正常"))
            logger.info("错误处理测试通过")
            return True
            
        except Exception as e:
            self.test_results.append(("错误处理测试", "失败", str(e)))
            logger.error(f"错误处理测试失败: {e}")
            return False
    
    async def test_完整流程验证(self):
        """测试完整的知识处理流程"""
        logger.info("开始完整流程验证")
        
        try:
            # 1. 上传测试文档
            document_id = await self.test_文档上传功能()
            if not document_id:
                raise Exception("文档上传失败，无法继续完整流程测试")
            
            # 2. 等待向量化完成
            await asyncio.sleep(2)
            
            # 3. 执行向量恢复
            vector_restore_success = await self.test_向量恢复功能()
            
            # 4. 执行知识搜索
            search_results = await self.test_知识搜索功能("测试文档")
            
            # 5. 验证完整流程结果
            if search_results and len(search_results) > 0:
                self.test_results.append(("完整流程测试", "通过", "上传->向量化->搜索流程正常"))
                logger.info("完整流程验证通过")
                return True
            else:
                self.test_results.append(("完整流程测试", "部分成功", "上传成功但搜索无结果"))
                logger.warning("完整流程验证部分成功：上传成功但搜索无结果")
                return False
                
        except Exception as e:
            self.test_results.append(("完整流程测试", "失败", str(e)))
            logger.error(f"完整流程验证失败: {e}")
            return False
    
    def 生成测试报告(self):
        """生成测试报告"""
        print("\n" + "="*60)
        print("知识搜索功能修复验证测试报告")
        print("="*60)
        
        total_tests = len(self.test_results)
        passed_tests = len([r for r in self.test_results if r[1] == "通过"])
        failed_tests = len([r for r in self.test_results if r[1] == "失败"])
        warning_tests = len([r for r in self.test_results if r[1] == "警告"])
        
        print(f"\n测试统计:")
        print(f"总测试数: {total_tests}")
        print(f"通过: {passed_tests}")
        print(f"失败: {failed_tests}")
        print(f"警告: {warning_tests}")
        
        print(f"\n详细结果:")
        for test_name, status, details in self.test_results:
            status_symbol = "✓" if status == "通过" else "✗" if status == "失败" else "⚠"
            print(f"{status_symbol} {test_name}: {status} - {details}")
        
        print("\n" + "="*60)
        
        # 返回测试是否整体成功
        return failed_tests == 0
    
    async def cleanup(self):
        """清理测试环境"""
        try:
            if self.db_manager:
                await self.db_manager.close()
            logger.info("测试环境清理完成")
        except Exception as e:
            logger.error(f"测试环境清理失败: {e}")

async def main():
    """主测试函数"""
    test_suite = 知识搜索功能修复验证测试()
    
    try:
        # 初始化测试环境
        if not await test_suite.setup():
            print("测试环境初始化失败，退出测试")
            return False
        
        # 执行各项测试
        await test_suite.test_系统初始化()
        await test_suite.test_数据库连接状态()
        await test_suite.test_向量恢复功能()
        await test_suite.test_错误处理机制()
        await test_suite.test_完整流程验证()
        
        # 生成测试报告
        success = test_suite.生成测试报告()
        
        return success
        
    except Exception as e:
        logger.error(f"测试执行过程中发生异常: {e}")
        return False
        
    finally:
        # 清理测试环境
        await test_suite.cleanup()

if __name__ == "__main__":
    # 运行测试
    success = asyncio.run(main())
    
    if success:
        print("\n🎉 所有测试通过！知识搜索功能修复验证成功。")
        sys.exit(0)
    else:
        print("\n❌ 部分测试失败，请检查源码并修复问题。")
        sys.exit(1)